import geopandas as gpd
import pyvista as pv
from pyproj import Transformer, CRS
import os
import logging
from pathlib import Path
import sys
import subprocess
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading

# 确保tqdm正确安装并导入
try:
    from tqdm import tqdm
except ImportError:
    print("自动安装tqdm进度条库...")
    subprocess.check_call([sys.executable, "-m", "pip", "install", "tqdm"])
    from tqdm import tqdm

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 线程安全的网格列表（用锁保护）
thread_safe_meshes = []
meshes_lock = threading.Lock()

def get_crs_from_prj(prj_path):
    """从.prj文件获取坐标参考系统"""
    try:
        encodings = ['utf-8', 'gbk', 'gb2312', 'latin-1']
        for encoding in encodings:
            try:
                with open(prj_path, 'r', encoding=encoding) as f:
                    prj_content = f.read()
                return CRS.from_wkt(prj_content)
            except UnicodeDecodeError:
                continue
        logger.warning(f"所有编码都无法解析.prj文件")
        return None
    except Exception as e:
        logger.warning(f"无法从.prj文件读取坐标系统: {str(e)}")
        return None

def generate_single_mesh(row, height_field, default_height):
    """多线程生成单个要素的3D网格"""
    single_meshes = []
    try:
        polygons = row["converted_coords"]
        if not polygons:
            return single_meshes

        # 处理高度
        if height_field:
            height = row[height_field]
            try:
                height = float(height)
                if height <= 0:
                    height = default_height
            except (ValueError, TypeError):
                height = default_height
        else:
            height = default_height

        # 处理多边形
        for poly_idx, polygon in enumerate(polygons):
            if len(polygon) < 3:
                logger.warning(f"跳过顶点不足的多边形（要素{row.name}，子多边形{poly_idx}）")
                continue

            # 构建顶点和面
            bottom_vertices = [(x, y, 0.0) for x, y in polygon]
            top_vertices = [(x, y, height) for x, y in polygon]
            all_vertices = bottom_vertices + top_vertices
            num_vertices = len(all_vertices)

            n = len(polygon)
            faces = []
            # 底面三角形化
            for i in range(1, n - 1):
                faces.extend([3, 0, i, i + 1])
            # 顶面三角形化
            for i in range(1, n - 1):
                faces.extend([3, n, n + i, n + i + 1])
            # 侧面四边形
            for i in range(n - 1):
                faces.extend([4, i, i + 1, n + i + 1, n + i])
            faces.extend([4, n - 1, 0, n, n + n - 1])  # 闭合侧面

            # 创建网格并添加属性
            mesh = pv.PolyData(all_vertices, faces)
            feature_id = row.get('ID', row.name)
            mesh["feature_id"] = [feature_id] * num_vertices
            mesh["height"] = [height] * num_vertices
            mesh["name"] = [f"feature_{feature_id}_{poly_idx}"] * num_vertices

            single_meshes.append(mesh)
    except Exception as e:
        logger.error(f"生成要素{row.name}的网格失败: {str(e)}")
    return single_meshes

def merge_chunk(chunk):
    """合并单个分块的网格（供多线程调用）"""
    if not chunk:
        return None
    merged = chunk[0]
    for mesh in chunk[1:]:
        merged = merged.merge(mesh)
    return merged

def shapefile_to_unity_model(shapefile_dir, output_path=None, height_field=None, default_height=10,
                             max_workers=4, merge_chunks=4):
    """
    分块多线程合并优化版
    :param merge_chunks: 合并分块数（建议与CPU核心数相同）
    """
    try:
        # 1. 查找Shapefile
        shp_files = list(Path(shapefile_dir).glob("*.shp"))
        if not shp_files:
            raise FileNotFoundError("未找到.shp文件")
        shp_path = shp_files[0]
        logger.info(f"找到目标Shapefile: {shp_path}")

        # 2. 读取数据（多编码尝试）
        logger.info("读取Shapefile数据...")
        encodings_to_try = ['gbk', 'gb2312', 'utf-8', 'latin-1']
        gdf = None
        for encoding in encodings_to_try:
            try:
                gdf = gpd.read_file(shp_path, encoding=encoding)
                logger.info(f"使用编码 {encoding} 成功读取数据")
                break
            except UnicodeDecodeError:
                continue
        if gdf is None:
            raise ValueError("所有编码均无法读取数据")

        logger.info(f"成功读取 {len(gdf)} 个要素")

        # 3. 坐标转换
        src_crs = gdf.crs
        prj_path = shp_path.with_suffix('.prj')
        if src_crs is None and prj_path.exists():
            src_crs = get_crs_from_prj(prj_path)

        transformer = None
        if src_crs and "GK_CM_129E" in src_crs.name:
            try:
                transformer = Transformer.from_crs(src_crs, "EPSG:4326", always_xy=True)
                logger.info("坐标将转换为WGS84经纬度")
            except Exception as e:
                logger.warning(f"坐标转换失败: {str(e)}")

        # 4. 转换坐标
        logger.info("批量转换坐标...")
        def convert_coords(geom):
            if geom.geom_type not in ['Polygon', 'MultiPolygon']:
                return None
            polygons = [geom] if geom.geom_type == 'Polygon' else list(geom.geoms)
            converted = []
            for poly in polygons:
                coords = list(poly.exterior.coords)
                if transformer:
                    try:
                        converted_coords = [transformer.transform(x, y) for x, y in coords]
                        converted.append(converted_coords)
                    except:
                        converted.append(coords)
                else:
                    converted.append(coords)
            return converted
        gdf["converted_coords"] = gdf["geometry"].apply(convert_coords)
        valid_gdf = gdf[gdf["converted_coords"].notna()].copy()
        logger.info(f"有效要素数: {len(valid_gdf)}")

        # 5. 多线程生成网格
        logger.info(f"多线程生成网格（{max_workers}线程）...")
        global thread_safe_meshes
        thread_safe_meshes = []
        futures = []
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            for _, row in valid_gdf.iterrows():
                future = executor.submit(generate_single_mesh, row, height_field, default_height)
                futures.append(future)

            for future in tqdm(as_completed(futures), total=len(futures), desc="生成网格"):
                result = future.result()
                if result:
                    with meshes_lock:
                        thread_safe_meshes.extend(result)

        if not thread_safe_meshes:
            raise ValueError("未生成任何网格")
        logger.info(f"共生成 {len(thread_safe_meshes)} 个网格")

        # 6. 分块多线程合并（核心优化）
        logger.info(f"分块多线程合并（{merge_chunks}个分块）...")

        # 切分成多个子块
        chunk_size = max(1, len(thread_safe_meshes) // merge_chunks)
        chunks = [thread_safe_meshes[i:i+chunk_size] for i in range(0, len(thread_safe_meshes), chunk_size)]

        # 多线程合并每个子块
        chunk_results = []
        with ThreadPoolExecutor(max_workers=merge_chunks) as executor:
            futures = [executor.submit(merge_chunk, chunk) for chunk in chunks]
            for future in tqdm(as_completed(futures), total=len(futures), desc="合并子块"):
                result = future.result()
                if result:
                    chunk_results.append(result)

        # 最终合并所有子块结果
        logger.info("合并子块结果...")
        combined_mesh = chunk_results[0]
        for mesh in tqdm(chunk_results[1:], total=len(chunk_results)-1, desc="最终合并"):
            combined_mesh = combined_mesh.merge(mesh)

        # 7. 导出模型
        if not output_path:
            output_dir = shp_path.parent
            output_name = f"{shp_path.stem}_unity_model.fbx"
            output_path = os.path.join(output_dir, output_name)
        else:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)

        logger.info(f"导出模型到: {output_path}")
        combined_mesh.save(output_path)
        logger.info("模型导出成功！")
        return output_path

    except Exception as e:
        logger.error(f"转换失败: {str(e)}", exc_info=True)
        raise

if __name__ == "__main__":
    # -------------------------- 配置区 --------------------------
    SHAPEFILE_DIR = r"C:\Users\ma\Desktop\28单位林班面"
    HEIGHT_FIELD = None  # 可尝试 "MIAN_JI"（面积字段）
    DEFAULT_HEIGHT = 10
    MAX_WORKERS = 4  # 生成网格的线程数
    MERGE_CHUNKS = 4  # 合并分块数（建议等于CPU核心数）
    OUTPUT_PATH = None
    # ------------------------------------------------------------

    try:
        result_path = shapefile_to_unity_model(
            shapefile_dir=SHAPEFILE_DIR,
            output_path=OUTPUT_PATH,
            height_field=HEIGHT_FIELD,
            default_height=DEFAULT_HEIGHT,
            max_workers=MAX_WORKERS,
            merge_chunks=MERGE_CHUNKS
        )
        print(f"\n✅ 转换完成！模型路径: {result_path}")
    except Exception as e:
        print(f"\n❌ 转换失败: {str(e)}")
